from __future__ import division
import numpy as np
from Constant_Physics import Physics
from Tools import *
from MRTools import *
from src.ShallowDonor import Shallow_Donor
from src.SDRmodel import *
import matplotlib.pyplot as plt
import matplotlib.mlab as mlab
import os

class Analysis():

    def __init__(self, clk, donor,fortran_compiled=True):
        self.clk=clk
        self.fortran_compiled_functions=fortran_compiled
        self.donor=donor
        
        self.sX={}
        self.sY={}
        self.SX={}
        self.SY={}
        
        self.kX={}
        self.kY={}
        self.KX={}
        self.KY={}
        
        self.hX={}
        self.hY={}
        self.HX={}
        self.HY={}
        
        self.extension=".csv"
        self.newline="\n"
        self.delimiter="\t"
        self.home_path=os.getcwd()#Current path to this program
        self.data_path=self.home_path+os.sep+"data"+os.sep#For saving later on data
        self.noise=0
        self.noise_psd=0

    def load_data(self,label,file_name):
        self.sX[label]=[]
        self.sY[label]=[]
        f = open(self.data_path+file_name+self.extension, 'r')
        res=f.readlines()
        for line in res:
            X,Y=line.split(',')
            self.sX[label].append(float(X))
            self.sY[label].append(float(Y))
        f.close()
        self.sX[label]=np.array(self.sX[label])
        self.sY[label]=np.array(self.sY[label])
        
        self.SX[label]=np.fft.fftshift(np.fft.fftfreq(self.sX[label].size, (self.sX[label][-1]-self.sX[label][0])/self.sX[label].size))
        self.SY[label]=np.fft.fft(np.array(self.sY[label]))
        
    def load_noise(self,label,file_name):
        self.noiseX[label]=[]
        self.noiseY[label]=[]
        f = open(self.data_path+file_name+self.extension, 'r')
        res=f.readlines()
        for line in res:
            X,Y=line.split(',')
            self.noiseX[label].append(float(X))
            self.noiseY[label].append(float(Y))
        f.close()
        self.noiseX[label]=np.array(self.noiseX[label])
        self.noiseY[label]=np.array(self.noiseY[label])
        Fs=len(self.noiseX[label])/(self.noiseX[label][-1]-self.noiseX[label][0])
        self.noiseY_psd[label],self.noiseX_psd[label]=mlab.psd(np.array(self.noiseY[label]),Fs=Fs)
        
    def calculate_response(self,label,sigma,mu,amp):
        self.kX[label]=self.sX[label]
        self.kY[label]=-self.DDnormal(np.array(self.kX[label]),sigma,mu,amp)
        self.KX[label]=np.fft.fftshift(np.fft.fftfreq(self.kX[label].size, (self.sX[label][-1]-self.sX[label][0])/self.sX[label].size))
        self.KY[label]=np.fft.fft(self.kY[label])

    def Deconvolve(self,label):
        self.HX[label]=self.KX[label]
        self.HY[label]=np.zeros([self.SY[label].size])
        for i in range(self.SY[label].size):
            self.HY[label][i]=self.SY[label][i]/self.KY[label][i]
        self.hY[label]=np.fft.fftshift(np.fft.ifft(self.HY[label]))

    def normal(self,x,sigma,mu,amp):
        return amp/(np.sqrt(2*np.pi)*sigma)*np.exp(-(x-mu)**2/(2*sigma**2))
    
    def Dnormal(self,x,sigma,mu,amp):
        return self.normal(x,sigma,mu,amp)*(-(x-mu)/sigma**2)
    
    def DDnormal(self,x,sigma,mu,amp):
        return self.normal(x,sigma,mu,amp)*(-1/sigma**2) + self.Dnormal(x,sigma,mu,amp)*(-(x-mu)/sigma**2)

if __name__ == "__main__":
    clk=Computation_time()
    Bi = Shallow_Donor(clk,"Bi")
    Bvec=Bfield(clk,0.066056)
    an=Analysis(clk,Bi)
    
    label="20-1"
    an.load_data(label,"20-1")
    an.calculate_response(label,0.37,65.369,7e-3)
       
    an.Deconvolve(label)
    
    #print an.HY[label]
    plt.plot(an.hY[label])
    #plt.plot(sig.deconvolve(an.kY[label],an.kY[label]))
    plt.show()
    
    